package com.flyone.fuxiao.controller;

import com.flyone.fuxiao.common.BaseResponse;
import com.flyone.fuxiao.common.ErrorCode;
import com.flyone.fuxiao.common.ResultUtils;
import com.flyone.fuxiao.model.dto.file.GenerateImgRequest;
import com.flyone.fuxiao.model.dto.img2img.Img2imgRequest;
import com.flyone.fuxiao.model.dto.txt2img.OverrideSettings;
import com.flyone.fuxiao.model.dto.txt2img.Txt2ImgRequest;
import com.flyone.fuxiao.model.entity.User;
import com.flyone.fuxiao.model.entity.sd.Model;
import com.flyone.fuxiao.model.vo.StarFireResponseVO;
import com.flyone.fuxiao.model.vo.ImgResponseVO;
import com.flyone.fuxiao.service.PicturePostService;
import com.flyone.fuxiao.service.UserService;
import io.swagger.annotations.Api;
import io.swagger.annotations.ApiOperation;
import org.apache.commons.lang3.StringUtils;
import org.springframework.web.bind.annotation.*;
import org.springframework.web.multipart.MultipartFile;

import javax.annotation.Resource;
import javax.servlet.http.HttpServletRequest;
import java.io.IOException;
import java.util.ArrayList;
import java.util.Base64;
import java.util.List;

/**
 * 模型调用 接口
 *
 * @author flyone
 * @date 2024/3/9
 */
@RestController
@RequestMapping("/model")
@Api(tags = "模型调用接口")
public class ModelController {

    @Resource
    private PicturePostService picturePostService;
    @Resource
    private UserService userService;

    @PostMapping("/generateImg")
    @ApiOperation("星火模型图片生成")
    public BaseResponse<StarFireResponseVO> generationImg(@RequestBody GenerateImgRequest generateImgRequest,
                                                          HttpServletRequest request,
                                                          @RequestParam("token") String token) {
        User loginUser = userService.getLoginUser(request);
        // 调用模型服务生成图片
        StarFireResponseVO responseVO = picturePostService.generationImg(generateImgRequest, token, loginUser);
        return ResultUtils.success(responseVO);
    }

    @PostMapping("/txt2img")
    @ApiOperation("Stable diffusion模型文字生成图片")
    public BaseResponse<ImgResponseVO> txt2img(@RequestBody Txt2ImgRequest txt2ImgRequest, HttpServletRequest request,
                                               @RequestParam("token") String token) {
        if (txt2ImgRequest == null || StringUtils.isBlank(txt2ImgRequest.getPrompt())) {
            return ResultUtils.error(ErrorCode.PARAMS_ERROR.getCode(), ErrorCode.PARAMS_ERROR.getMessage());
        }
        User loginUser = userService.getLoginUser(request);
        ImgResponseVO imgResponseVO = picturePostService.txt2img(txt2ImgRequest, token, loginUser);
        return ResultUtils.success(imgResponseVO);
    }

//    @PostMapping("/img2img")
//    @ApiOperation("Stable diffusion模型图片生成图片")
//    public BaseResponse<String> img2img(@RequestPart("images") List<MultipartFile> multipartFiles,
//                                               HttpServletRequest request,
//                                               @RequestParam("type") String type) throws IOException {
//        User loginUser = userService.getLoginUser(request);
//        List<String> images = new ArrayList<>();
//        if (CollectionUtils.isNotEmpty(multipartFiles)) {
//            for (MultipartFile multipartFile : multipartFiles) {
//                String imgBase64 = Base64.getEncoder().encodeToString(multipartFile.getBytes());
//                images.add(imgBase64);
//            }
//        }
//        Img2imgRequest img2imgRequest = new Img2imgRequest();
//        img2imgRequest.setInit_images(images);
//        OverrideSettings overrideSettings = new OverrideSettings();
//        overrideSettings.setSd_model_checkpoint("majicmix7");
//        img2imgRequest.setOverride_settings(overrideSettings);
//        String url = picturePostService.img2img(img2imgRequest, type, loginUser);
//        return ResultUtils.success(url);
//    }

    @PostMapping("/img2img")
    @ApiOperation("Stable diffusion模型图片生成图片")
    public BaseResponse<String> img2img(@RequestPart("images") List<MultipartFile> multipartFiles,
                                        HttpServletRequest request,
                                        @RequestParam("type") String type) throws IOException {
        User loginUser = userService.getLoginUser(request);
        List<String> images = new ArrayList<>();

        // 判断type参数是否为"handDraw"
        if ("handDraw".equals(type)) {
            // 手绘生图逻辑
            for (MultipartFile multipartFile : multipartFiles) {
                String imgBase64 = Base64.getEncoder().encodeToString(multipartFile.getBytes());
                images.add(imgBase64);
            }
            Img2imgRequest img2imgRequest = new Img2imgRequest();
            img2imgRequest.setInit_images(images);
            OverrideSettings overrideSettings = new OverrideSettings();
            overrideSettings.setSd_model_checkpoint("majicmix7");
            img2imgRequest.setOverride_settings(overrideSettings);
            String url = picturePostService.img2img(img2imgRequest, "handDraw", loginUser);
            return ResultUtils.success(url);
        } else {
            // 图生图逻辑
            for (MultipartFile multipartFile : multipartFiles) {
                String imgBase64 = Base64.getEncoder().encodeToString(multipartFile.getBytes());
                images.add(imgBase64);
            }
            Img2imgRequest img2imgRequest = new Img2imgRequest();
            img2imgRequest.setInit_images(images);
            OverrideSettings overrideSettings = new OverrideSettings();
            overrideSettings.setSd_model_checkpoint("majicmix7");
            img2imgRequest.setOverride_settings(overrideSettings);
            String url = picturePostService.img2img(img2imgRequest, "img2img", loginUser);
            return ResultUtils.success(url);
        }
    }


    @GetMapping("/sd-models")
    @ApiOperation("Stable diffusion模型获取")
    public BaseResponse<Model[]> getSdModels() {
        Model[] models = picturePostService.getSdModels();
        return ResultUtils.success(models);
    }

//    @GetMapping("/history")
//    @ApiOperation("历史记录获取")
//    public BaseResponse<List<PicturePost>> getSdHistory() {
//        List<String> history = picturePostService.getSdHistory();
//        return ResultUtils.success(history);
//    }

}
